#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# @Author: renjin@bit.edu.cn
# @Date  : 2024-09-12


"""
【节点名称】：
    SpireViewGroundingDINONode
【依赖项安装】：
    pip install spirems
    git clone https://gitee.com/jario-jin/GroundingDINO.git
    cd GroundingDINO/
    pip3 install -e . -i https://pypi.tuna.tsinghua.edu.cn/simple/
【订阅类型】：
    spirecv_msgs::GroundingDINOJob （输入GroundingDINO任务）
【发布类型】：
    spirecv_msgs::GroundingDINOResult （GroundingDINO任务结果）
【构造参数说明】：
    parameter_file (str): 全局参数文件
【节点参数】：
    groundingdino_py (str): GroundingDINO模型配置文件
    groundingdino_pth (str): GroundingDINO模型参数文件
    box_threshold (float)
    text_threshold (float)
【备注】：
    无
"""

import os

import cv2
# if using Apple MPS, fall back to CPU for unsupported ops
# os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
import numpy as np
import torch
import matplotlib.pyplot as plt
from PIL import Image
import threading
from queue import Queue
import time
from spirems import Subscriber, Publisher, cvimg2sms, sms2cvimg, def_msg
from BaseNode import BaseNode
import base64
from groundingdino.util.inference import load_model, predict, annotate
import groundingdino.datasets.transforms as T


def groundingdino_input(image: np.ndarray):
    transform = T.Compose(
        [
            T.RandomResize([800], max_size=1333),
            T.ToTensor(),
            T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
        ]
    )
    image_source = Image.fromarray(image)  # RGB image
    image_transformed, _ = transform(image_source, None)
    return image, image_transformed


class GroundingDINOPipeline(threading.Thread):
    def __init__(self, groundingdino_model, client_id, box_threshold, text_threshold, ip, port):
        threading.Thread.__init__(self)
        self.groundingdino_model = groundingdino_model
        self.box_threshold = box_threshold
        self.text_threshold = text_threshold
        self.client_id = client_id
        self.msg_queue = Queue()
        self._result_writer = Publisher(
            '/{}/OutputGroundingDINOResult'.format(self.client_id), 'spirecv_msgs::GroundingDINOResult',
            ip = ip, port = port
        )
        self.is_running = True
        self.image_buff = {}
        self.start()

    def put_msg(self, msg):
        self.msg_queue.put(msg)

    def quit(self):
        print('quit')
        self._result_writer.kill()
        self.is_running = False
        self.msg_queue.put(None)

    def run(self):
        while self.is_running:
            msg = self.msg_queue.get(block=True)
            if msg is None:
                break

            try:
                set_image = 0
                if 'type' in msg['image'] and msg['image']['type'] == 'sensor_msgs::CompressedImage':
                    # t1 = time.time()
                    img_bgr = sms2cvimg(msg['image'])
                    img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
                    self.image_buff.clear()
                    self.image_buff[msg["image_id"]] = img_rgb
                    set_image = 1
                    # print('dt1: {}'.format(time.time() - t1))

                input_text = None
                if len(msg['input_text']):
                    input_text = msg['input_text']

                if msg["image_id"] in self.image_buff:
                    # t2 = time.time()
                    res_msg = def_msg('spirecv_msgs::GroundingDINOResult')
                    res_msg["set_image"] = set_image
                    res_msg["image_id"] = msg["image_id"]
                    if "input_id" in msg:
                        res_msg["input_id"] = msg["input_id"]
                    if input_text is None:
                        res_msg['detections'] = []
                    else:
                        image_source, image = groundingdino_input(self.image_buff[msg["image_id"]])
                        boxes, logits, phrases = predict(
                            model=self.groundingdino_model,
                            image=image,
                            caption=input_text,
                            box_threshold=self.box_threshold,
                            text_threshold=self.text_threshold
                        )
                        results = {
                            'boxes': boxes.cpu().numpy().astype(np.float64),
                            'logits': logits.cpu().numpy().astype(np.float64),
                            'phrases': phrases
                        }
                        detections = []
                        for box, score, phrase in zip(results['boxes'], results['logits'], results['phrases']):
                            bbox = box.tolist()
                            bbox.append(score)
                            detection = {
                                'bbox': bbox,
                                'category': phrase
                            }
                            detections.append(detection)
                        res_msg['detections'] = detections
                        # print(res_msg)
                        # res_msg["results"] = results
                    self._result_writer.publish(res_msg)
            except Exception as e:
                print('Client: {}, Error: {}'.format(self.client_id, e))
                self.is_running = False


class SpireViewGroundingDINONode(threading.Thread, BaseNode):
    def __init__(
        self,
        ip: str = '127.0.0.1',
        port: int = 9094,
        parameter_file: str = ''
    ):
        threading.Thread.__init__(self)
        BaseNode.__init__(self, self.__class__.__name__, 'SpireViewGroundingDINO', ip=ip, port=port,
                          parameter_file=parameter_file)

        self.job_queue = Queue()
        self.ip = ip
        self.port = port
        self.queue_pool.append(self.job_queue)

        self._image_reader = Subscriber(
            '/SpireView/InputGroundingDINOJob', 'spirecv_msgs::GroundingDINOJob', self.job_callback,
            ip=ip, port=port
        )

        self.groundingdino_py = self.get_param("groundingdino_py", "G:/deep/GroundingDINO/groundingdino/config/GroundingDINO_SwinB_cfg.py")
        self.groundingdino_pth = self.get_param("groundingdino_pth", "G:/deep/GroundingDINO/groundingdino_swinb_cogcoor.pth")
        self.groundingdino_model = load_model(self.groundingdino_py, self.groundingdino_pth)

        self.box_threshold = self.get_param("box_threshold", 0.35)
        self.text_threshold = self.get_param("text_threshold", 0.25)
        self.client_buff = {}
        self.start()

    def release(self):
        BaseNode.release(self)
        self._image_reader.kill()

    def job_callback(self, msg):
        self.job_queue.put(msg)

    def run(self):
        while self.is_running():
            msg = self.job_queue.get(block=True)
            if msg is None:
                break
            t1 = time.time()

            if 'client_id' not in msg or not isinstance(msg['client_id'], str):
                continue

            if msg['client_id'] not in self.client_buff:
                pipeline = GroundingDINOPipeline(
                    self.groundingdino_model, msg['client_id'], self.box_threshold, self.text_threshold,
                    self.ip, self.port
                )
                self.client_buff[msg['client_id']] = {
                    'pipeline': pipeline,
                    'time': time.time()
                }
                to_del = []
                for _key, _value in self.client_buff.items():
                    if time.time() - _value['time'] > 600:
                        to_del.append(_key)
                for _key in to_del:
                    self.client_buff[_key]['pipeline'].quit()
                    del self.client_buff[_key]
            else:
                self.client_buff[msg['client_id']]['time'] = time.time()

            if self.client_buff[msg['client_id']]['pipeline'].is_running:
                self.client_buff[msg['client_id']]['pipeline'].put_msg(msg)
            else:
                del self.client_buff[msg['client_id']]
        self.release()
        print('{} quit!'.format(self.__class__.__name__))


if __name__ == '__main__':
    import uuid
    parameter_file = 'default_params.json'
    dino = SpireViewGroundingDINONode(parameter_file=parameter_file)
    """
    img = cv2.imread('G:/deep/segment-anything-2/notebooks/images/truck.jpg')
    pub = Publisher('/SpireView/InputGroundingDINOJob', 'spirecv_msgs::GroundingDINOJob')
    sam_msg = def_msg('spirecv_msgs::GroundingDINOJob')
    sam_msg['client_id'] = str(uuid.uuid4()).replace('-', '_')
    sam_msg['image_id'] = str(uuid.uuid4()).replace('-', '_')
    sam_msg['image'] = cvimg2sms(img)

    while 1:
        sam_msg['image'] = cvimg2sms(img)
        sam_msg['input_text'] = ""
        pub.publish(sam_msg, True)
        sam_msg['image'] = {}
        for i in range(10):
            time.sleep(1)
            sam_msg['input_text'] = "track"
            pub.publish(sam_msg, True)
        time.sleep(30)
    """
